package com.yx.vectordb.impl;

import cn.hutool.core.bean.BeanUtil;
import com.yx.vectordb.api.IMetaData;
import com.yx.vectordb.api.IVectorDB;
import com.yx.vectordb.impl.lucene.LuceneSearchResult;
import com.yx.vectordb.impl.lucene.LuceneVectorEngine;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.sql.*;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
public class YxVectorDB<T extends IMetaData> implements IVectorDB<T> {

    private Connection connection;
    private int dimension;
    private File dbDirectory;
    private File dbPath;
    private File indexPath;
    private String jdbcUrl;
    private Class<T> tClass;
    MetaTable table;

    LuceneVectorEngine luceneVectorEngine;

    private Map<Class,MetaTable> metaTableMap = new HashMap<>();

    public void open(File dbDirectory, Class<T> metaInfoBeanClass, int dimension) throws Exception{
        this.tClass = metaInfoBeanClass;
        this.dbDirectory = dbDirectory;
        try{
            this.dbPath = new File(this.dbDirectory,"meta.db");
            jdbcUrl = "jdbc:h2:" + dbPath.getAbsolutePath();
            connection = getConnection();
        }catch (Exception exception){
            throw exception;
        }
        try{
            loadIndexInfo(metaInfoBeanClass);
        }catch (Exception ex){
            this.indexPath = new File(this.dbDirectory,"vector.store");
            this.dimension = dimension;
            if (!existTable("CONFIG")) {
                String createConfigSQL = "CREATE TABLE CONFIG (" +
                        "K CHARACTER VARYING PRIMARY KEY," +
                        "V CHARACTER VARYING" +
                        ")";
                executeSQL(createConfigSQL);
            }
            setValue("vector.filePath",this.indexPath.getAbsolutePath());
            setValue("vector.dimension",this.dimension+"");
            table = MetaTable.fromBean(metaInfoBeanClass);
            table.setName("META_INFO");
            if(!existTable(table.getName().toUpperCase())){
                createTable(table);
            }
            try {
                luceneVectorEngine = new LuceneVectorEngine(this.indexPath.getAbsolutePath(),this.dimension);
            } catch (Exception e) {
                throw e;
            }
        }
    }

    private void loadIndexInfo(Class<T> metaInfoBeanClass) throws Exception {
        this.indexPath = new File(this.dbDirectory,getValue("vector.filePath"));
        this.dimension = Integer.parseInt(getValue("vector.dimension"));
        table = MetaTable.fromBean(metaInfoBeanClass);
        table.setName("META_INFO");
        if(!existTable(table.getName().toUpperCase())){
            createTable(table);
        }
        try {
             luceneVectorEngine = new LuceneVectorEngine(this.indexPath.getAbsolutePath(),this.dimension);
        } catch (Exception e) {
            throw e;
        }

    }

    public String getValue(String key) throws SQLException {
        String sql = "select V from CONFIG where K=?";
        try (PreparedStatement ps = connection.prepareStatement(sql)) {
            ps.setString(1, key);
            try (ResultSet rs = ps.executeQuery()) {
                if (rs.next()) {
                    return rs.getString("V");
                }
            }
        }
        return null;
    }

    public String setValue(String key, String value) throws SQLException {
        String oldValue = getValue(key);
        if (oldValue != null) {
            // 更新已存在的值
            String updateSql = "UPDATE CONFIG SET V=? WHERE K=?";
            try (PreparedStatement ps = connection.prepareStatement(updateSql)) {
                ps.setString(1, value);
                ps.setString(2, key);
                ps.executeUpdate();
            }
        } else {
            // 插入新值
            String insertSql = "INSERT INTO CONFIG (K, V) VALUES (?, ?)";
            try (PreparedStatement ps = connection.prepareStatement(insertSql)) {
                ps.setString(1, key);
                ps.setString(2, value);
                ps.executeUpdate();
            }
        }
        return oldValue;
    }


    private boolean existTable(String name) throws SQLException {
        String sql = "SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = 'PUBLIC' AND TABLE_NAME = ?";
        try (Connection conn = getConnection();
             PreparedStatement stmt = conn.prepareStatement(sql)) {
            stmt.setString(1, name.toUpperCase());
            try (ResultSet rs = stmt.executeQuery()) {
                if (rs.next()) {
                    return rs.getInt(1) > 0;
                }
            }
        }
        return false;
    }
    
    private void createTable(MetaTable table) throws SQLException {
        StringBuilder sb = new StringBuilder();
        sb.append("create table PUBLIC.").append(table.getName()).append("\n(\n");
        sb.append("\tID   CHARACTER VARYING not null,").append("\n");
        for(Column column: table.getColumnList()){
            sb.append("\t").append(column.getName()).append(" ").append(column.getType()).append(",\n");
        }
        sb.append("\tconstraint "+table.getName().toUpperCase()+"_PK"+" primary key (ID)");
        sb.append("\n)");
        String createTableSQL = sb.toString();

        List<String> createIndexSQLList = new ArrayList<>();
        for(Column column: table.getColumnList()){
            String indexSQL = "create index PUBLIC."+table.getName().toUpperCase()+"_"+column.getName().toUpperCase()+"_INDEX on PUBLIC."+table.getName().toUpperCase()+" ("+column.getName().toUpperCase()+")";
            createIndexSQLList.add(indexSQL);
        }
        List<String> all = new ArrayList<>();
        all.add(createTableSQL);
        all.addAll(createIndexSQLList);
        executeSQL(all);
    }

    public void executeSQL(List<String> sqls) throws SQLException {
        if (sqls == null || sqls.isEmpty()) {
            return;
        }
        executeSQL(sqls.toArray(new String[0]));
    }

    public void executeSQL(String... sqlList) throws SQLException {
        Connection connection = null;
        boolean originalAutoCommit = true;
        try {
            connection = getConnection();
            // 保存原始的自动提交设置
            originalAutoCommit = connection.getAutoCommit();
            // 开启事务
            connection.setAutoCommit(false);
            
            for(String sql : sqlList) {
                try (Statement stmt = connection.createStatement()) {
                    stmt.execute(sql);
                }
            }
            
            // 提交事务
            connection.commit();
        } catch (Exception ex) {
            // 发生异常时回滚事务
            if (connection != null) {
                try {
                    connection.rollback();
                } catch (SQLException e) {
                    log.error("事务回滚失败", e);
                }
            }
            throw ex;
        } finally {
            if (connection != null) {
                try {
                    // 恢复原始的自动提交设置
                    connection.setAutoCommit(originalAutoCommit);
                } catch (SQLException e) {
                    log.error("恢复自动提交设置失败", e);
                }
                close(connection);
            }
        }
    }

    public void executeSQL(String sql) throws SQLException {
        Connection connection = null;
        try{
            connection = getConnection();
            try (Statement stmt = connection.createStatement()) {
                stmt.execute(sql);
            }
        }catch (Exception ex){
            throw ex;
        }finally {
            if(connection!=null){
                close(connection);
            }
        }
    }

    public Connection getConnection() throws SQLException {
        return DriverManager.getConnection(jdbcUrl);
    }

    public void close(Connection connection){
        try{
            connection.close();
        }catch (Exception ex){
            log.error(ex.getMessage(),ex);
        }
    }


    public String addVector(String id, float[] vector, T meta) throws Exception {
        try {
            if(id!=null){
                // 检查是否存在旧的元数据
                T existingMeta = getMetaById(id);
                if (existingMeta != null) {
                    // 如果存在旧的向量ID，先删除旧的向量
                    luceneVectorEngine.deleteVector(id);
                }
            }else {
                id = UUID.randomUUID().toString();
            }

            // 添加新的向量
            luceneVectorEngine.addVector(id,vector);
            meta.setId(id);
            // 保存或更新元数据
            insertOrUpdateMeta(meta);
            return id;
        } catch (Exception e) {
            throw e;
        }
    }

    @Override
    public void deleteVector(String id) throws Exception {
        try{
            luceneVectorEngine.deleteVector(id);
        }catch (Exception ex){}

        try{
            deleteMetaByIds(id);
        }catch (Exception ex){}

    }

    @Override
    public void deleteVector(Collection<String> ids) throws Exception {
        try{
            ids.forEach(id-> {
                try {
                    luceneVectorEngine.deleteVector(id);
                } catch (IOException e) {
                }
            });
        }catch (Exception ex){}

        try{
            deleteMetaByIds(ids);
        }catch (Exception ex){}
    }

    @Override
    public void deleteBy(String whereCondition) throws Exception {
        // 如果whereCondition是blank，则直接退出，防止删除所有
        if (whereCondition == null || whereCondition.trim().isEmpty()) {
            log.warn("删除条件为空，为防止误删除所有数据，操作已取消");
            return;
        }
        
        // 查询符合条件的所有ID
        List<String> idsToDelete = new ArrayList<>();
        String sql = "SELECT ID FROM " + table.getName() + " WHERE " + whereCondition;
        
        try (PreparedStatement ps = connection.prepareStatement(sql);
             ResultSet rs = ps.executeQuery()) {
            
            while (rs.next()) {
                String id = rs.getString("ID");
                if (id != null && !id.trim().isEmpty()) {
                    idsToDelete.add(id);
                }
            }
            
            // 如果没有找到匹配的记录，直接返回
            if (idsToDelete.isEmpty()) {
                log.info("没有找到符合条件的记录: {}", whereCondition);
                return;
            }
            
            // 使用已有的方法删除向量和元数据
            log.info("根据条件 '{}' 删除 {} 条记录", whereCondition, idsToDelete.size());
            deleteVector(idsToDelete);
            
        } catch (SQLException e) {
            log.error("根据条件删除记录失败: " + whereCondition, e);
            throw new SQLException("根据条件删除记录失败", e);
        }
    }

    @Override
    public List<VectorSearchResult<T>> search(float[] queryVector, int k, String whereCondition) throws Exception {
        return select(queryVector,0f,k,whereCondition);
    }

    @Override
    public List<VectorSearchResult<T>> search(float[] queryVector, float threshold, String whereCondition) throws Exception {
        return select(queryVector,threshold,100,whereCondition);
    }

    @Override
    public List<VectorSearchResult<T>> search(float[] queryVector, float threshold, int k, String whereCondition) throws Exception {
        return select(queryVector,threshold,k,whereCondition);
    }

    private List<VectorSearchResult<T>> select(float[] vector, float minScore, int k, String where) throws Exception {
        try {
            // 1. 使用Faiss进行向量搜索
            List<LuceneSearchResult> luceneSearchResultList = luceneVectorEngine.search(vector,k);

            luceneSearchResultList = luceneSearchResultList.stream().filter(sr->sr.getScore()>minScore).collect(Collectors.toList());

            String[] vectorIds = luceneSearchResultList.stream().map(sr->sr.getId()).collect(Collectors.toList()).toArray(new String[0]);
            Float[] distances = luceneSearchResultList.stream().map(luceneSearchResult -> luceneSearchResult.getScore()).collect(Collectors.toList()).toArray(new Float[0]);
            
            // 如果没有找到相似向量，直接返回空列表
            if (luceneSearchResultList.isEmpty()) {
                return Collections.emptyList();
            }
            
            // 2. 构建SQL查询条件
            StringBuilder sql = new StringBuilder();
            sql.append("SELECT * FROM ").append(table.getName());
            sql.append(" WHERE ID IN (");
            for (int i = 0; i < vectorIds.length; i++) {
                if (i > 0) sql.append(",");
                sql.append("\'"+vectorIds[i]+"\'");
            }
            sql.append(")");
            
            // 如果有额外的where条件，添加到查询中
            if (where != null && !where.trim().isEmpty()) {
                sql.append(" AND (").append(where).append(")");
            }
            
            // 3. 查询元数据
            Map<String, T> metaMap = new HashMap<>();
            try (PreparedStatement ps = connection.prepareStatement(sql.toString());
                 ResultSet rs = ps.executeQuery()) {
                
                while (rs.next()) {
                    T instance = tClass.getDeclaredConstructor().newInstance();
                    ResultSetMetaData metaData = rs.getMetaData();
                    int columnCount = metaData.getColumnCount();
                    Map<String,Object> beanMap = new HashMap<>();
                    for (int i = 1; i <= columnCount; i++) {
                        String columnName = metaData.getColumnName(i);
                        Object value = rs.getObject(i);
                        beanMap.put(columnName,value);
                    }
                    BeanUtil.copyProperties(beanMap,instance,true);
                    
                    metaMap.put(instance.getId(), instance);
                }
            }
            
            // 4. 按照向量搜索结果的顺序创建Rank列表
            List<VectorSearchResult<T>> results = new ArrayList<>();
            for (int i = 0; i < vectorIds.length; i++) {
                T meta = metaMap.get(vectorIds[i]);
                if (meta != null) {  // 只添加满足where条件的结果
                    results.add(new VectorSearchResult<T>(meta, distances[i]));
                }
            }
            
            return results;
        } catch (Exception e) {
            log.error("向量搜索失败", e);
            throw new SQLException("向量搜索失败", e);
        }
    }

    private void insertOrUpdateMeta(T meta) throws SQLException {
        T existingMeta = getMetaById(meta.getId());
        if (existingMeta != null) {
            // 更新已存在的记录
            StringBuilder updateSql = new StringBuilder();
            updateSql.append("UPDATE ").append(table.getName()).append(" SET ");
            List<String> updateColumns = new ArrayList<>();
            
            for (Column column : table.getColumnList()) {
                String columnName = column.getName().toUpperCase();
                Object value = BeanUtil.getFieldValue(meta, column.getName().toLowerCase());
                if (value != null) {
                    updateColumns.add(columnName + "=?");
                }
            }
            
            updateSql.append(String.join(",", updateColumns));
            updateSql.append(" WHERE ID=?");
            
            try (PreparedStatement ps = connection.prepareStatement(updateSql.toString())) {
                int paramIndex = 1;
                for (Column column : table.getColumnList()) {
                    Object value = BeanUtil.getFieldValue(meta, column.getName().toLowerCase());
                    if (value != null) {
                        ps.setObject(paramIndex++, value);
                    }
                }
                ps.setString(paramIndex, meta.getId());
                ps.executeUpdate();
            }
        } else {
            // 插入新记录
            StringBuilder insertSql = new StringBuilder();
            insertSql.append("INSERT INTO ").append(table.getName()).append(" (ID,");
            List<String> columns = new ArrayList<>();
            List<String> values = new ArrayList<>();
            
            for (Column column : table.getColumnList()) {
                columns.add(column.getName().toUpperCase());
                values.add("?");
            }
            
            insertSql.append(String.join(",", columns)).append(") VALUES (?," + String.join(",", values) + ")");
            
            try (PreparedStatement ps = connection.prepareStatement(insertSql.toString())) {
                ps.setString(1, meta.getId());
                int paramIndex = 2;
                for (Column column : table.getColumnList()) {
                    Object value = BeanUtil.getFieldValue(meta, column.getName());
                    ps.setObject(paramIndex++, value);
                }
                ps.executeUpdate();
            }
        }
    }

    private T getMetaById(String id) throws SQLException {
        List<T> results = selectMeta("ID = '" + id + "'");
        return results.isEmpty() ? null : results.get(0);
    }

    private void deleteMetaByIds(String id) throws SQLException {
        if (id == null || id.trim().isEmpty()) {
            throw new SQLException("删除元数据失败：ID不能为空");
        }
        
        String sql = "DELETE FROM " + table.getName() + " WHERE ID = ?";
        try (PreparedStatement ps = connection.prepareStatement(sql)) {
            ps.setString(1, id);
            int affectedRows = ps.executeUpdate();
            if (affectedRows == 0) {
                log.warn("没有找到ID为 '{}' 的元数据记录", id);
            } else {
                log.debug("成功删除ID为 '{}' 的元数据记录", id);
            }
        } catch (SQLException e) {
            log.error("删除元数据记录失败，ID: " + id, e);
            throw new SQLException("删除元数据记录失败", e);
        }
    }

    private void deleteMetaByIds(Collection<String> ids) throws SQLException {
        if (ids == null || ids.isEmpty()) {
            return; // 如果集合为空，直接返回
        }
        
        // 构建IN子句的参数占位符
        StringBuilder placeholders = new StringBuilder();
        for (int i = 0; i < ids.size(); i++) {
            if (i > 0) {
                placeholders.append(",");
            }
            placeholders.append("?");
        }
        
        String sql = "DELETE FROM " + table.getName() + " WHERE ID IN (" + placeholders + ")";
        
        try (PreparedStatement ps = connection.prepareStatement(sql)) {
            int paramIndex = 1;
            for (String id : ids) {
                if (id != null && !id.trim().isEmpty()) {
                    ps.setString(paramIndex++, id);
                }
            }
            
            int affectedRows = ps.executeUpdate();
            if (affectedRows == 0) {
                log.warn("没有找到任何匹配的元数据记录进行删除");
            } else {
                log.debug("成功删除 {} 条元数据记录", affectedRows);
            }
        } catch (SQLException e) {
            log.error("批量删除元数据记录失败", e);
            throw new SQLException("批量删除元数据记录失败", e);
        }
    }

    private List<T> selectMeta(String where) throws SQLException {
        List<T> result = new ArrayList<>();
        String sql = "SELECT * FROM " + table.getName();
        if (where != null && !where.trim().isEmpty()) {
            sql += " WHERE " + where;
        }
        
        try (PreparedStatement ps = connection.prepareStatement(sql);
             ResultSet rs = ps.executeQuery()) {
            
            while (rs.next()) {
                T instance = tClass.getDeclaredConstructor().newInstance();
                ResultSetMetaData metaData = rs.getMetaData();
                int columnCount = metaData.getColumnCount();
                Map<String,Object> beanMap = new HashMap<>();
                for (int i = 1; i <= columnCount; i++) {
                    String columnName = metaData.getColumnName(i);
                    Object value = rs.getObject(i);
                    beanMap.put(columnName,value);
                    // 忽略大小写的属性复制

                }
                BeanUtil.copyProperties(beanMap,instance);
                result.add(instance);
            }
        } catch (Exception e) {
            throw new SQLException("查询元数据失败", e);
        }
        
        return result;
    }

    public void close(){
        try{
            connection.close();
        }catch (Exception ex){
        }

        try{
            luceneVectorEngine.close();
        }catch (Exception ex){
        }
    }

}